/*
 * @Date: 2022-02-12 12:04:58
 * @LastEditors: tich425
 * @LastEditTime: 2022-02-12 12:04:58
 * @FilePath: \TodoList\middleware\limit.go
 * @Description: 限流
 */
package middleware

import (
	"time"
	"todo-list/global"
	"todo-list/res"

	"github.com/gin-gonic/gin"
)

/**
 * @description: 按次数单路由限流
 * @param {int} limit
 */
func LimitRoute(limit int64) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := c.ClientIP() + c.FullPath()
		// 读取请求次数
		value, err := global.REDIS_DB.Get(key).Int64()
		if err != nil {
			// 初始化请求次数
			global.REDIS_DB.Set(key, 1, time.Hour)
			c.Next()
		} else if value < limit {
			// 请求次数 +1
			global.REDIS_DB.Incr(key)
			c.Next()
		} else {
			// 请求次数超过限制
			res.FailWithMessage(res.TooManyRequests, "请求过于频繁, 请一小时后重试", c)
			c.Abort()
		}
	}
}

/**
 * @description: 按时间单路由限流
 * @param {int} limit
 */
func LimitAverageRoute(limit int64) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := c.ClientIP() + c.FullPath()
		// 读取值并验证
		if err := global.REDIS_DB.Get(key).Err(); err != nil {
			// 初始化请求次数
			global.REDIS_DB.Set(key, 1, time.Hour/time.Duration(limit))
			c.Next()
		} else {
			res.FailWithMessage(res.TooManyRequests, "请求过于频繁, 请稍后尝试", c)
			c.Abort()
		}
	}
}
